package uk.co.jgo2rdf.tools;

/*
 * Copyright: This program is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU Lesser General Public License (GNU-LGPL) as published by
 *     the Free Software Foundation, either version 3 of the License, or
 *     (at your option) any later version.
 *
 *     This program is distributed in the hope that it will be useful,
 *     but WITHOUT ANY WARRANTY; without even the implied warranty of
 *     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *     GNU Lesser General Public License for more details.
 *
 *     You should have received a copy of the GNU Lesser General Public License
 *     along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.*;
import java.util.Map.Entry;
/**
 * A one-directional indexing class for the NCBI taxonomy 
 * 
 * @author lysenkoa
 *
 */
public class NCBICommonAncestor {
	private final Map<Integer,Set<SpeciesInt>> speciesIndex = new HashMap<Integer,Set<SpeciesInt>>();
	
	public static final void main(String [] args) throws Exception{
		NCBICommonAncestor nc = new NCBICommonAncestor();
		nc.parse(new FileInputStream("/home/lysenkoa/Downloads/nodes.dmp"));
		System.err.println("Yay! Enough memory!");
		System.err.println(nc.getCommonAncestor(27706, 7955)); // the right answer is 186625
		
	}
	
	/**
	 * Loads a taxonomy file
	 * @param file - nodes.dmp from the NCBI ftp
	 * @throws Exception
	 */
	public void parse(InputStream is) throws Exception{
		final Map<Integer,Set<Integer>> index = new HashMap<Integer,Set<Integer>>();
		BufferedReader br = new BufferedReader(new InputStreamReader(is));
		String line;
			while((line = br.readLine()) != null){
				String [] data = line.split("\\|");
				Integer child = Integer.valueOf(data[0].trim());
				Integer parent = Integer.valueOf(data[1].trim());
				Set<Integer> set = index.get(child);
				if(set == null){
					set = new HashSet<Integer>();
					index.put(child, set);
				}
				set.add(parent);
			}

		br.close();
		

		
		int additions = 0;
		do{
			additions = 0;
			Set<Integer> add = new HashSet<Integer>();
			for(Entry<Integer,Set<Integer>> ent : index.entrySet()){
				int size = ent.getValue().size();
				for(Integer s : ent.getValue()){
					add.addAll(index.get(s));
				}
				ent.getValue().addAll(add);
				add.clear();
				size = ent.getValue().size() - size;
				additions = additions + size;
			}
		}while(additions > 0);
		
		Map<Integer, SpeciesInt> map = new HashMap<Integer, SpeciesInt>();
		
		for(Entry<Integer, Set<Integer>> ent : index.entrySet()){
			map.put(ent.getKey(), new SpeciesInt(ent.getKey(), ent.getValue().size()));
		}
		
		for(Entry<Integer, Set<Integer>> ent : index.entrySet()){
			speciesIndex.put(ent.getKey(), convert(ent.getValue(), map));
			ent.getValue().clear();
		}
	}
	
	private Set<SpeciesInt> convert(Set<Integer> set, Map<Integer, SpeciesInt> map){
		Set<SpeciesInt> result = new HashSet<SpeciesInt>();
		for(Integer i : set){
			result.add(map.get(i));
		}
		return result;
	}
	
	
	
	private static class SpeciesInt implements Comparable<SpeciesInt>{
		public final Integer id;
		public final Integer order;

		public SpeciesInt(Integer id, Integer order) {
			super();
			this.id = id;
			this.order = order;
		}
		
		@Override
		public int hashCode(){
			return id.hashCode();
		}
		
		@Override
		public boolean equals(Object obj){
			if(obj instanceof SpeciesInt){
				return id.equals(((SpeciesInt)obj).id);
			}
			return false;
		}

		@Override
		public int compareTo(SpeciesInt arg) {
			return this.order.compareTo(arg.order);
		}

		@Override
		public String toString(){
			return id.toString();
		}
	}
	
	/**
	 * Get the lowest common taxa for a pair of species A and B
	 * @param speciesA
	 * @param speciesB
	 * 
	 * @return ncbi taxid
	 */
	public Integer getCommonAncestor(Integer speciesA, Integer speciesB){
		if(speciesA.equals(speciesB)){
			return speciesB;
		}
		Set<SpeciesInt> a = new HashSet<SpeciesInt>(speciesIndex.get(speciesA));
		a.retainAll(speciesIndex.get(speciesB));
		return Collections.max(a).id;
	}
}
